'''
Created on Aug 11, 2009
@author: mjbommar
'''

from tfaces.collection.Experiment import *
from tfaces.collection.util import *
import pickle


class GraphExperiment(Experiment):
    '''
    The Graph Experiment class loads an Experiment that produces graph output.
    
    The output should be of the following form:
    
    parameter1=value
    parameter2=value
    .
    .
    .
    # tail,head[,weight]
    0,1,1
    0,1,3
    .
    .
    .
    
    
    Any lines that begin with # that do not fall directly before the data section will be treated as comments.
    
    '''
    
    droneParameters = ['reportFileName', 'reportFrequency', 'outputDirName', 'runNumber', 'initialParametersFileName', 'initialAgentsFileName', 'saveRunEndState', 'reportFileNameSuffix']
    
    def __init__(self, experimentDirectory='.', reportFilePattern='report.[0-9]+'):
        '''
        The Drone Experiment constructor takes as arguments the directory in which the
        experiment was stored, and whether to parse the headers from the report file.
        '''

        # Call parent constructor
        Experiment.__init__(self, experimentDirectory, reportFilePattern)
        
        self.processGraphRuns(self.findRuns())
    
    def __repr__(self):
        '''
        Return a string representation.
        '''
        return 'tfaces.collection.GraphExperiment: %d parameters, %d runs' % (len(self.parameterKeys), len(self.runs))
    
        
    def processGraphRuns(self, runFiles):
        """
        This method processes a list of text run files.
        """
        runs = [self.processGraphRun(runFile) for runFile in runFiles]
        self.addRuns(runs)

    def processGraphRun(self, runFile):
        """
        Parse each report file and return a Run object.
        """
        # Initialize the variables for the Run object
        parameters = {}
        outputs = {}
        seed = 0

        # Load the data and split it into lines
        data = open(runFile).read()
        dataLines = [line.strip() for line in data.strip().splitlines()]
        seenHeader = False
        edges = []
        edgeWeights = []

        for lineNo, line in enumerate(dataLines):
            # Skip header lines
            if line.startswith('#') or line.startswith('@'):
                continue
            
            # Look for parameter definitions
            if '=' in line:
                name, value = [token.strip() for token in line.split('=')]
                if not isOneOf(value, [str, int, float]):
                    raise Exception("GraphExperiment::processTextRun: invalid definition type for %s on line %d" % (name, lineNo))
                
                if name == 'seed':
                    seed = int(value)
                elif name not in self.droneParameters:
                    parameters[name] = toNaturalType(value)

            elif len(line) > 0:
                # Now parse the data for each line and make sure it matches the header
                dataTokens = self.lineSplit(line)
                if len(dataTokens) not in [2,3]:
                    raise Exception("GraphExperiment::processTextRun: number of data columns not equal to number of headers on line %d" % lineNo)

                if len(dataTokens) == 3:
                    edges.append((int(dataTokens[0]), int(dataTokens[1])))
                    edgeWeights.append(float(dataTokens[2]))
                else:
                    edges.append(int(dataTokens[0]), int(dataTokens[1]))
            
        # Produce the actual output dictionary by forcing the data from the file into its natural type
        graph = igraph.Graph(edges)
        
        if len(edgeWeights) == len(edges):
            for i,e in enumerate(graph.es):
                e["weight"] = edgeWeights
        
        return GraphRun(parameters, graph, seed, False)
    
    def initDB(self):
        '''
        Initialize the database
        '''
        
        # Generate the parameter set table from the parameter list.
        sqlParameterSet = 'CREATE TABLE IF NOT EXISTS parameterset ('
        for parameter in self.parameterKeys:
            sqlParameterSet += '%s TEXT,' % parameter
        sqlParameterSet += 'id INTEGER PRIMARY KEY AUTOINCREMENT)'
        
        self.db.execute(sqlParameterSet)
        
        # Generate the run table from the output list.
        sqlRuns = 'CREATE TABLE IF NOT EXISTS runs (graph TEXT, seed INTEGER NOT NULL, parametersetid INTEGER NOT NULL)'
        
        self.db.execute(sqlRuns)
        self.db.commit()
    
    def getParameterSetID(self, parameters):
        '''
        Determine the parameter set ID from the given dictionary.
        '''
        
        # Initialize the SELECT and INSERT strings
        selectString = "SELECT id FROM parameterset WHERE "
        insertString = "INSERT INTO parameterset VALUES ("
        
        # Generate the SELECT and INSERT strings from the parameters
        for i, (k, v) in enumerate(parameters.iteritems()):
            if i < len(parameters) - 1:
                selectString += '%s = "%s" AND ' % (k, v)
                insertString += '"%s",' % v
            else:
                selectString += '%s = "%s"' % (k, v)
                insertString += '"%s", NULL)' % v
        
        ''' Run the SELECT string, see if the row exists.  If it exists, return the ID.
        Otherwise, INSERT the row, then return the new ID.'''
        result = self.db.execute(selectString).fetchone()
        if result:
            return result[0]
        else:
            self.db.execute(insertString)
            return self.db.execute(selectString).fetchone()[0]
    
    def save(self, fileName):
        '''
        Write the experiment out to a SQLite database file.
        '''
        # If the experiment file already exists, raise an exception.
        if os.path.exists(fileName):
            raise Exception("Experiment::save: %s already exists" % (fileName))
        
        # Initialize the database.
        self.db = sqlite3.connect(fileName)
        self.initDB()
        
        sqlInsertData = []
        
        # Generate the SQL queries for the data for all runs.  
        for run in self.runs:
            ''' Get the parameter set ID for the given run.  Then, we
            generate the template for the INSERT statement.''' 
            parameterSetID = self.getParameterSetID(run.parameters)
            baseInsertString = 'INSERT INTO runs (graph, seed, parametersetid) VALUES (?,?,?)'
            
            ''' Now generate the data tuples that we pass to the template INSERT
            statement in batch mode. '''
            sqlInsertData.append((pickle.dumps(run.graph), run.seed, parameterSetID))

        # Finally add the data and commit.
        self.db.executemany(baseInsertString, sqlInsertData)
        self.numrows = len(sqlInsertData)
        self.db.commit()
        self.db.close()
    
